{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Use Vertex Session Service for Free with Agent Engine, ADK, and Vertex Express Mode\n",
        "\n",
        "This tutorial extends from the [Quickstart example](https://google.github.io/adk-docs/get-started/quickstart/) for [Agent Development Kit](https://google.github.io/adk-docs/get-started/).\n",
        "\n",
        "We'll embark on building a **Weather Bot agent**, creating single agent that can look up weather, we will connect it with some free Vertex services like the VertexAiSessionService, allowing for users to save their sessions on vertex!\n",
        "\n",
        "**What is ADK Again?**\n",
        "\n",
        "As a reminder, ADK is a Python framework designed to streamline the development of applications powered by Large Language Models (LLMs). It offers robust building blocks for creating agents that can reason, plan, utilize tools, interact dynamically with users, and collaborate effectively within a team.\n",
        "\n",
        "**In this tutorial, you will master:**\n",
        "\n",
        "*   ✅ **Tool Definition & Usage:** Crafting Python functions (`tools`) that grant agents specific abilities (like fetching data) and instructing agents on how to use them effectively.\n",
        "*   ✅ **Agent Engine APIs:** Deploying local agents and using agent engine capabilities like the VertexAiSession service for free\n",
        "\n",
        "**End State Expectation:**\n",
        "\n",
        "By completing this tutorial, you will have built a functional weather agent that can utilize the agent engine building blocks.\n",
        "\n",
        "**Prerequisites:**\n",
        "\n",
        "*   ✅ **Solid understanding of Python programming.**\n",
        "*   ✅ **Familiarity with Large Language Models (LLMs), APIs, and the concept of agents.**\n",
        "*   ❗ **Crucially: Completion of the ADK Quickstart tutorial(s) or equivalent foundational knowledge of ADK basics (Agent, Runner, SessionService, basic Tool usage).** This tutorial builds directly upon those concepts.\n",
        "*   ✅ **API Keys** for the LLMs you intend to use (e.g., Google AI Studio for Gemini, OpenAI Platform, Anthropic Console).\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "**Ready to build your agent team? Let's dive in!**"
      ],
      "metadata": {
        "id": "TpDkHI2BNeay"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Setup and Installation\n",
        "# Install ADK\n",
        "\n",
        "!pip install google-adk -q\n",
        "\n",
        "print(\"Installation complete.\")"
      ],
      "metadata": {
        "id": "RsED19m8K3M6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sbwxKypOSBkN"
      },
      "outputs": [],
      "source": [
        "# @title Import necessary libraries\n",
        "import os\n",
        "import asyncio\n",
        "from google import adk\n",
        "from google.adk.agents import Agent\n",
        "from google.adk.sessions import VertexAiSessionService\n",
        "from google.adk.memory import VertexAiMemoryBankService\n",
        "from google.adk.runners import Runner\n",
        "from google.genai import types # For creating message Content/Parts\n",
        "\n",
        "import warnings\n",
        "# Ignore all warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configure API Keys (Replace with your actual keys!)\n",
        "We can use Vertex Services like VertexAiSessionService and access models for free through Vertex Express Mode! Sign up with your google account here: https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/overview to unlock free access to certain services and models for free without the need of adding your credit card.\n",
        "\n",
        "Vertex Express mode combined with ADK allow for the creation of advanced agents for free! You will have access to certain Gemini models and Agent Engine services like Session and Memory, all for free without a billing account!"
      ],
      "metadata": {
        "id": "NqEJDL_6k58J"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3mNsVI5eSDOi"
      },
      "outputs": [],
      "source": [
        "# Gemini API Key (Get from Vertex Express Mode)\n",
        "express_mode_api_key = \"YOUR-EXPRESS-MODE-API-KEY\" #@param {type:\"string\"}\n",
        "os.environ[\"GOOGLE_API_KEY\"] = express_mode_api_key\n",
        "# Set vertex to true\n",
        "os.environ[\"GOOGLE_GENAI_USE_VERTEXAI\"] = \"True\"\n",
        "\n",
        "# --- Verify Keys (Optional Check) ---\n",
        "print(\"API Keys Set:\")\n",
        "print(f\"Google API Key set: {'Yes' if os.environ.get('GOOGLE_API_KEY') and os.environ['GOOGLE_API_KEY'] != 'INSERT API KEY HERE' else 'No (REPLACE PLACEHOLDER!)'}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "When creating an agent, we need to choose the model we want. The Vertex Express mode API key allows for the use of several gemini models for free, and you can use any of the models listed here: https://cloud.google.com/vertex-ai/generative-ai/docs/start/express-mode/overview#models"
      ],
      "metadata": {
        "id": "6-xwIAk-lf1T"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MI_qvZJrSJuR"
      },
      "outputs": [],
      "source": [
        "# --- Define Model Constants for easier use ---\n",
        "\n",
        "# Use an allowlisted model for EasyGCP, we will use gemini 2.5\n",
        "MODEL = \"gemini-2.5-flash-lite\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F7LZM3ysSOMu"
      },
      "source": [
        "---\n",
        "\n",
        "## Step 1: Your First Agent \\- Basic Weather Lookup\n",
        "\n",
        "Let's begin by building the fundamental component of our Weather Bot: a single agent capable of performing a specific task – looking up weather information. This involves creating two core pieces:\n",
        "\n",
        "1. **A Tool:** A Python function that equips the agent with the *ability* to fetch weather data.  \n",
        "2. **An Agent:** The AI \"brain\" that understands the user's request, knows it has a weather tool, and decides when and how to use it.\n",
        "\n",
        "---\n",
        "\n",
        "**1\\. Define the Tool (`get_weather`)**\n",
        "\n",
        "In ADK, **Tools** are the building blocks that give agents concrete capabilities beyond just text generation. They are typically regular Python functions that perform specific actions, like calling an API, querying a database, or performing calculations.\n",
        "\n",
        "Our first tool will provide a *mock* weather report. This allows us to focus on the agent structure without needing external API keys yet. Later, you could easily swap this mock function with one that calls a real weather service.\n",
        "\n",
        "**Key Concept: Docstrings are Crucial\\!** The agent's LLM relies heavily on the function's **docstring** to understand:\n",
        "\n",
        "* *What* the tool does.  \n",
        "* *When* to use it.  \n",
        "* *What arguments* it requires (`city: str`).  \n",
        "* *What information* it returns.\n",
        "\n",
        "**Best Practice:** Write clear, descriptive, and accurate docstrings for your tools. This is essential for the LLM to use the tool correctly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ILy7YTCbSRAT"
      },
      "outputs": [],
      "source": [
        "# @title Define the get_weather Tool\n",
        "def get_weather(city: str) -> dict:\n",
        "    \"\"\"Retrieves the current weather report for a specified city.\n",
        "\n",
        "    Args:\n",
        "        city (str): The name of the city (e.g., \"New York\", \"London\", \"Tokyo\").\n",
        "\n",
        "    Returns:\n",
        "        dict: A dictionary containing the weather information.\n",
        "              Includes a 'status' key ('success' or 'error').\n",
        "              If 'success', includes a 'report' key with weather details.\n",
        "              If 'error', includes an 'error_message' key.\n",
        "    \"\"\"\n",
        "    print(f\"--- Tool: get_weather called for city: {city} ---\") # Log tool execution\n",
        "    city_normalized = city.lower().replace(\" \", \"\") # Basic normalization\n",
        "\n",
        "    # Mock weather data\n",
        "    mock_weather_db = {\n",
        "        \"newyork\": {\"status\": \"success\", \"report\": \"The weather in New York is sunny with a temperature of 25°C.\"},\n",
        "        \"london\": {\"status\": \"success\", \"report\": \"It's cloudy in London with a temperature of 15°C.\"},\n",
        "        \"tokyo\": {\"status\": \"success\", \"report\": \"Tokyo is experiencing light rain and a temperature of 18°C.\"},\n",
        "    }\n",
        "\n",
        "    if city_normalized in mock_weather_db:\n",
        "        return mock_weather_db[city_normalized]\n",
        "    else:\n",
        "        return {\"status\": \"error\", \"error_message\": f\"Sorry, I don't have weather information for '{city}'.\"}\n",
        "\n",
        "# Example tool usage (optional test)\n",
        "print(get_weather(\"New York\"))\n",
        "print(get_weather(\"Paris\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hAM0BqGWSTo5"
      },
      "source": [
        "---\n",
        "\n",
        "**2\\. Define the Agent (`weather_agent`)**\n",
        "\n",
        "Now, let's create the **Agent** itself. An `Agent` in ADK orchestrates the interaction between the user, the LLM, and the available tools.\n",
        "\n",
        "We configure it with several key parameters:\n",
        "\n",
        "* `name`: A unique identifier for this agent (e.g., \"weather\\_agent\\_v1\").  \n",
        "* `model`: Specifies which LLM to use (e.g., `MODEL_GEMINI_2_0_FLASH`). We'll start with a specific Gemini model.  \n",
        "* `description`: A concise summary of the agent's overall purpose. This becomes crucial later when other agents need to decide whether to delegate tasks to *this* agent.  \n",
        "* `instruction`: Detailed guidance for the LLM on how to behave, its persona, its goals, and specifically *how and when* to utilize its assigned `tools`.  \n",
        "* `tools`: A list containing the actual Python tool functions the agent is allowed to use (e.g., `[get_weather]`).\n",
        "\n",
        "**Best Practice:** Provide clear and specific `instruction` prompts. The more detailed the instructions, the better the LLM can understand its role and how to use its tools effectively. Be explicit about error handling if needed.\n",
        "\n",
        "**Best Practice:** Choose descriptive `name` and `description` values. These are used internally by ADK and are vital for features like automatic delegation (covered later)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Ho1COmKSUeV"
      },
      "outputs": [],
      "source": [
        "# @title Define the Weather Agent\n",
        "\n",
        "weather_agent = Agent(\n",
        "    name=\"weather_agent_v1\",\n",
        "    model=MODEL,\n",
        "    description=\"Provides weather information for specific cities.\",\n",
        "    instruction=\"You are a helpful weather assistant. \"\n",
        "                \"When the user asks for the weather in a specific city, \"\n",
        "                \"use the 'get_weather' tool to find the information. \"\n",
        "                \"If the tool returns an error, inform the user politely. \"\n",
        "                \"If the tool is successful, present the weather report clearly.\"\n",
        "                \"You can help the user find what city has better weather based on their preferences using the preload memory tool.\",\n",
        "    tools=[get_weather, adk.tools.preload_memory_tool.PreloadMemoryTool()], # Pass the function directly\n",
        ")\n",
        "\n",
        "print(f\"Agent '{weather_agent.name}' created using model '{MODEL}'.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "**3\\. Define the Agent Engine**\n",
        "\n",
        "An Agent Engine is a set of services that enables developers to deploy, manage, and scale AI agents in production. Agent Engine handles the infrastructure to scale agents in production so you can focus on creating applications. In Vertex Express mode, we have access to certain Agent Engine services for free, mainly the Session and Memory services, which allow for context management. Each session and memory is associated with an Agent Engine.\n",
        "\n",
        "We configure our Agent Engine with several key parameters:\n",
        "\n",
        "* `displayName`: A unique identifier for this agent engine (e.g., \"weather\\_agent\\_v1\").  \n",
        "* `description`: A concise summary of the agent engine's overall purpose. This can help you remember what is does."
      ],
      "metadata": {
        "id": "J3-rNFdr1urU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Create the Agent Engine\n",
        "import vertexai\n",
        "\n",
        "client = vertexai.Client(api_key=express_mode_api_key)\n",
        "agent_engine = client.agent_engines.create(config={\"display_name\": \"Test Agent Engine\", \"description\": \"My first Agent Engine\"})"
      ],
      "metadata": {
        "id": "goqjFK2SMNa4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "APP_NAME=agent_engine.api_resource.name\n",
        "APP_NAME"
      ],
      "metadata": {
        "id": "Ik0W51d_MW5Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "APP_ID=APP_NAME.split('/')[-1]\n",
        "APP_ID"
      ],
      "metadata": {
        "id": "_DSD83dghQpJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dvz7LDhbSZxL"
      },
      "source": [
        "---\n",
        "\n",
        "**4\\. Setup Runner and Session Service**\n",
        "\n",
        "To manage conversations and execute the agent, we need two more components:\n",
        "\n",
        "* `SessionService`: Responsible for managing conversation history and state for different users and sessions. The `VertexAiSessionService` is an implementation that stores everything in vertex, allowing for persistent session storage. It keeps track of the messages exchanged. We'll explore state persistence more in Step 4\\.  \n",
        "* `Runner`: The engine that orchestrates the interaction flow. It takes user input, routes it to the appropriate agent, manages calls to the LLM and tools based on the agent's logic, handles session updates via the `SessionService`, and yields events representing the progress of the interaction."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Create Our Initial Session\n",
        "\n",
        "# Create Vertex AI Session through ADK to use locally\n",
        "session_service = VertexAiSessionService(agent_engine_id=APP_ID)\n",
        "\n",
        "# Create Vertex AI Memory through ADK to use locally\n",
        "memory_service = VertexAiMemoryBankService(agent_engine_id=APP_ID)\n",
        "\n",
        "USER_ID = \"INSERT_USER_ID\" #@param {type:\"string\"}\n",
        "session = await session_service.create_session(app_name=APP_ID, user_id=USER_ID)\n",
        "SESSION_ID = session.id\n",
        "session"
      ],
      "metadata": {
        "id": "Nykf_acxMk5g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Create an Agent Runner\n",
        "\n",
        "# Connect with ADK. ADK will also use the easygcp key to generate content\n",
        "print(f\"Session created: App='{APP_ID}', User='{USER_ID}', Session='{SESSION_ID}'\")\n",
        "# --- Runner ---\n",
        "# Key Concept: Runner orchestrates the agent execution loop.\n",
        "runner = Runner(\n",
        "    agent=weather_agent, # The agent we want to run\n",
        "    app_name=APP_ID,   # Associates runs with our app\n",
        "    session_service=session_service, # Uses vertex session service\n",
        "    memory_service=memory_service # Uses vertex memory service\n",
        ")\n",
        "print(f\"Runner created for agent '{runner.agent.name}'.\")"
      ],
      "metadata": {
        "id": "_80NpAeF6-i6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5zKGVwRkSduA"
      },
      "source": [
        "---\n",
        "\n",
        "**5\\. Interact with the Agent**\n",
        "\n",
        "We need a way to send messages to our agent and receive its responses. Since LLM calls and tool executions can take time, ADK's `Runner` operates asynchronously.\n",
        "\n",
        "We'll define an `async` helper function (`call_agent_async`) that:\n",
        "\n",
        "1. Takes a user query string.  \n",
        "2. Packages it into the ADK `Content` format.  \n",
        "3. Calls `runner.run_async`, providing the user/session context and the new message.  \n",
        "4. Iterates through the **Events** yielded by the runner. Events represent steps in the agent's execution (e.g., tool call requested, tool result received, intermediate LLM thought, final response).  \n",
        "5. Identifies and prints the **final response** event using `event.is_final_response()`.\n",
        "\n",
        "**Why `async`?** Interactions with LLMs and potentially tools (like external APIs) are I/O-bound operations. Using `asyncio` allows the program to handle these operations efficiently without blocking execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yZJr8lbkSebH"
      },
      "outputs": [],
      "source": [
        "# @title Define Agent Interaction Function\n",
        "\n",
        "from google.genai import types # For creating message Content/Parts\n",
        "\n",
        "async def call_agent_async(query: str, runner, user_id, session_id):\n",
        "  \"\"\"Sends a query to the agent and prints the final response.\"\"\"\n",
        "  print(f\"\\n>>> User Query: {query}\")\n",
        "\n",
        "  # Prepare the user's message in ADK format\n",
        "  content = types.Content(role='user', parts=[types.Part(text=query)])\n",
        "\n",
        "  final_response_text = \"Agent did not produce a final response.\" # Default\n",
        "\n",
        "  # Key Concept: run_async executes the agent logic and yields Events.\n",
        "  # We iterate through events to find the final answer.\n",
        "  async for event in runner.run_async(user_id=user_id, session_id=session_id, new_message=content):\n",
        "      # You can uncomment the line below to see *all* events during execution\n",
        "      print(f\"  [Event] Author: {event.author}, Type: {type(event).__name__}, Final: {event.is_final_response()}, Content: {event.content}\")\n",
        "\n",
        "      # Key Concept: is_final_response() marks the concluding message for the turn.\n",
        "      if event.is_final_response():\n",
        "          if event.content and event.content.parts:\n",
        "             # Assuming text response in the first part\n",
        "             final_response_text = event.content.parts[0].text\n",
        "          elif event.actions and event.actions.escalate: # Handle potential errors/escalations\n",
        "             final_response_text = f\"Agent escalated: {event.error_message or 'No specific message.'}\"\n",
        "          # Add more checks here if needed (e.g., specific error codes)\n",
        "          break # Stop processing events once the final response is found\n",
        "\n",
        "  print(f\"<<< Agent Response: {final_response_text}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6DQSqrqk5ic"
      },
      "source": [
        "---\n",
        "\n",
        "**6\\. Run the Conversation**\n",
        "\n",
        "Finally, let's test our setup by sending a few queries to the agent. We wrap our `async` calls in a main `async` function and run it using `await`.\n",
        "\n",
        "Watch the output:\n",
        "\n",
        "* See the user queries.  \n",
        "* Notice the `--- Tool: get_weather called... ---` logs when the agent uses the tool.  \n",
        "* Observe the agent's final responses labelled with \"Agent Response\", including how it handles the case where weather data isn't available (for Paris)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mEd2QhHyUKY8"
      },
      "outputs": [],
      "source": [
        "# @title Run the Initial Conversation\n",
        "\n",
        "# We need an async function to await our interaction helper\n",
        "async def run_conversation():\n",
        "    await call_agent_async(\"What is the weather like in London?\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID)\n",
        "\n",
        "    await call_agent_async(\"How about Paris?\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID) # Expecting the tool's error message\n",
        "\n",
        "    await call_agent_async(\"Tell me the weather in New York\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID)\n",
        "    await call_agent_async(\"I prefer the weather in New York, that sounds nicer than the weather in London\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID)\n",
        "    await call_agent_async(\"What cities did I ask you about previously?\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID)\n",
        "\n",
        "# Execute the conversation using await in an async context (like Colab/Jupyter)\n",
        "await run_conversation()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xbUzAGvsmB2a"
      },
      "source": [
        "---\n",
        "\n",
        "Congratulations\\! You've successfully built and interacted with your first ADK agent, and used the Vertex Session Service for free!"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---\n",
        "\n",
        "**8\\. Test out Agent Memory**\n",
        "\n",
        "Lets see if the agent will remember our preferences from our previous session.\n",
        "- Here we can create a new session, then ask the agent about something we talked about in the previous session. In this example, we ask about our weather preferences after talking about preferring New York weather in the previous conversation.\n",
        "- The agent should utilize the vertex memory service to retrieve relevant details about the user, then utilize that in its responses."
      ],
      "metadata": {
        "id": "mdCneYx5tbcL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Create a Memory Based on the Previous Session\n",
        "\n",
        "# We can generate a memory given the previous session id\n",
        "session = await session_service.get_session(app_name=APP_ID, session_id=SESSION_ID, user_id = USER_ID)\n",
        "await memory_service.add_session_to_memory(session=session)"
      ],
      "metadata": {
        "id": "T1ZZua59gP4Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Test the Agent Memory\n",
        "\n",
        "# Create a new session, and lets see if it will remember our preferences based on our user id\n",
        "session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)\n",
        "SESSION_ID = session.id\n",
        "\n",
        "print(f\"New Session created: App='{APP_NAME}', User='{USER_ID}', Session='{SESSION_ID}'\")\n",
        "\n",
        "await call_agent_async(\"What city has the best weather for me based on my preferences?\",\n",
        "                                       runner=runner,\n",
        "                                       user_id=USER_ID,\n",
        "                                       session_id=SESSION_ID)"
      ],
      "metadata": {
        "id": "l6nw_N41eztw"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
